import torch.nn as nn
import torch
import torch.nn.functional as F
import torchvision.models as models


class SPPModule(nn.Module):
    def __init__(self, sizes, pool_type = "avg"):
        super().__init__()
        self.pools = nn.ModuleList()

        for size in sizes:
            if pool_type == "avg":
                self.pools.append(nn.AdaptiveAvgPool2d(size))
            elif pool_type == "max":
                self.pools.append(nn.AdaptiveMaxPool2d(size))

    def forward(self, x):
        resList = []
        for pool in self.pools:
            y = pool(x)
            y = torch.flatten(y,1)
            resList.append(y)
        res = torch.cat(resList, 1)
        return res




if __name__ == '__main__':
    pass



